package org.finesys.common.log.trace.executor;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;

/**
 * The purpose of this decorator is to capture and propagate the current request context and Mapped Diagnostic Context (MDC)
 * to a new thread when a task is submitted to an Executor for asynchronous execution.
 */
public class ContextCopyingDecorator implements TaskDecorator {

    /**
     * Decorates a Runnable to ensure that the current request context and MDC (Mapped Diagnostic Context)
     * are preserved and available in the new thread when the task is executed asynchronously.
     *
     * @param runnable The original Runnable to be decorated.
     * @return A new Runnable that wraps the original Runnable and ensures the context is preserved.
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        // Capture the current request attributes from the RequestContextHolder.
        RequestAttributes requestAttributes = RequestContextHolder.currentRequestAttributes();
        // Capture a copy of the current MDC context map.
        Map<String, String> map = MDC.getCopyOfContextMap();
        // Return a new Runnable that wraps the original Runnable.
        return () -> {
            try {
                // Set the captured request attributes in the new thread.
                RequestContextHolder.setRequestAttributes(requestAttributes);
                // Set the captured MDC context map in the new thread.
                if (map != null) {
                    MDC.setContextMap(map);
                }
                // Execute the original Runnable.
                runnable.run();
            } finally {
                // Reset the request attributes to avoid memory leaks.
                RequestContextHolder.resetRequestAttributes();
                // Clear the MDC to avoid memory leaks.
                MDC.clear();
            }
        };
    }
}
